// Package jwt
/**
* @Project : geektime-basic-go-study
* @File    : redis_jwt.go
* @IDE     : GoLand
* @Author  : Tvux
* @Date    : 2024/10/28 09:01
**/

package jwt

import (
	"errors"
	"fmt"
	"github.com/google/uuid"
	"strings"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/golang-jwt/jwt/v5"
	"github.com/redis/go-redis/v9"
)

var (
	AccessJWTKey256  = []byte("kYmG77wj31AhpiIrJcKomPZJ2zbT1fLvQ5LhNH5PkFjY8Cl7AUEsW8xaSxi3mdfteGQvcKVu1UJB0x9PS5vM2jX4GtppEoRo2GosGKsVxD61Xufujn7XeZcSEyx0hGcEl6fgsg3X9kGhEZJgztNPinVLDh2rYqeJsnYWDxKfbYf85I89KyOsXnJF9kGbrs8iMEfwK6Nx5CM8xaJ0jMoVLya1nPgYZYSOrOjeinhvPJJamFFvMfk2QhZeSdL3QwwG")
	RefreshJWTKey256 = []byte("l4T5r4GGN9eGHmZTKMuA0ppQu6x83WDXAPiaovlrkfNOrvmPUH5SA507MVed2ornIwnbdVIpkSx59VVlnaRyDSj8J3KEy29qOnkWYfPX42umll5RIJCaz5uf6ypEmzOFmo6SAETzmMkIVy48CijUW0a4UBmTyKdhCfdQ3YVXxlmjDPlzF5z4Jf7fXrLbRJnlnz7e8FvfI3RCBvTm8RVDOvgVo1Czz8M7EZIMu3LDKq01aOx1vM8HGwwTCxXPvPXl")
)

var (
	_ UserJWTHandler = (*RedisJWTHandler)(nil)
)

type RedisJWTHandler struct {
	redisClient redis.Cmdable
}

func NewRedisJWTHandler(redisClient redis.Cmdable) UserJWTHandler {
	return &RedisJWTHandler{
		redisClient: redisClient,
	}
}

func (Self *RedisJWTHandler) SetAccessAndRefreshJWT(ctx *gin.Context, uid int64) error {
	ssid := uuid.New().String()
	err := Self.setAccessJWT(ctx, uid, ssid)
	if err != nil {
		return err
	}
	err = Self.setRefreshJWT(ctx, uid, ssid)
	return err
}

func (Self *RedisJWTHandler) setAccessJWT(ctx *gin.Context, uid int64, ssid string) error {
	claims := UserClaims{
		RegisteredClaims: jwt.RegisteredClaims{
			ExpiresAt: jwt.NewNumericDate(time.Now().Add(time.Minute * 30)),
		},
		Uid:       uid,
		Ssid:      ssid,
		UserAgent: ctx.Request.UserAgent(),
	}
	token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
	tokenStr, err := token.SignedString(AccessJWTKey256)
	if err != nil {
		return err
	}

	ctx.Header("x-access-token", tokenStr)
	return nil
}

func (Self *RedisJWTHandler) setRefreshJWT(ctx *gin.Context, uid int64, ssid string) error {
	claims := RefreshCliams{
		RegisteredClaims: jwt.RegisteredClaims{
			ExpiresAt: jwt.NewNumericDate(time.Now().Add(time.Minute * 30)),
		},
		Uid:  uid,
		Ssid: ssid,
	}
	token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
	tokenStr, err := token.SignedString(RefreshJWTKey256)
	if err != nil {
		return err
	}

	ctx.Header("x-refresh-token", tokenStr)
	return nil
}

func (Self *RedisJWTHandler) RefreshAccessJWT(ctx *gin.Context) error {
	refreshJWT := Self.ExtractJWTFromHeader(ctx)
	refreshClaims, err := Self.parseRefrshJWT(refreshJWT)
	if err != nil {
		return fmt.Errorf("failed to parse refresh JWT: %w", err)
	}
	exists, err := Self.IsSessionExists(ctx, refreshClaims.Ssid)
	if err != nil {
		return fmt.Errorf("failed to check session existence: %w", err)
	}
	if exists {
		return errors.New("user has logout")
	}
	return Self.setAccessJWT(ctx, refreshClaims.Uid, refreshClaims.Ssid)
}

func (Self *RedisJWTHandler) parseRefrshJWT(refreshJWT string) (*RefreshCliams, error) {
	token, err := jwt.ParseWithClaims(refreshJWT, &RefreshCliams{}, func(token *jwt.Token) (interface{}, error) {
		if _, ok := token.Method.(*jwt.SigningMethodHMAC); !ok {
			return nil, fmt.Errorf("unexpected signing method: %v", token.Header["alg"])
		}
		return RefreshJWTKey256, nil
	})

	if err != nil {
		return nil, err
	}

	if !token.Valid {
		return nil, errors.New("invalid token")
	}

	if claims, ok := token.Claims.(*RefreshCliams); ok {
		return claims, nil
	} else {
		return nil, errors.New("failed to parse RefreshCliams")
	}
}

func (Self *RedisJWTHandler) ParseAccessJWT(accessJWT string) (*UserClaims, error) {
	token, err := jwt.ParseWithClaims(accessJWT, &UserClaims{}, func(token *jwt.Token) (interface{}, error) {
		if _, ok := token.Method.(*jwt.SigningMethodHMAC); !ok {
			return nil, fmt.Errorf("unexpected signing method: %v", token.Header["alg"])
		}
		return AccessJWTKey256, nil
	})

	if err != nil {
		return nil, err
	}

	if !token.Valid {
		return nil, errors.New("invalid token")
	}

	if claims, ok := token.Claims.(*UserClaims); ok {
		return claims, nil
	} else {
		return nil, errors.New("failed to parse UserClaims")
	}
}

func (Self *RedisJWTHandler) ExtractJWTFromHeader(ctx *gin.Context) string {
	authorization := ctx.GetHeader("Authorization")
	tokenStr := strings.TrimPrefix(authorization, "Bearer ")
	return tokenStr
}

func (Self *RedisJWTHandler) InvalidateJWT(ctx *gin.Context) error {
	ctx.Header("x-access-token", "")
	ctx.Header("x-refresh-token", "")

	claims, _ := ctx.Get("UserClaims")
	userClaims, ok := claims.(*UserClaims)
	if !ok {
		return fmt.Errorf("")
	}
	return Self.redisClient.Set(ctx, Self.keyFormat(userClaims.Ssid), "", time.Hour*24*7).Err()
}

func (Self *RedisJWTHandler) IsSessionExists(ctx *gin.Context, ssid string) (bool, error) {
	exists, err := Self.redisClient.Exists(ctx, Self.keyFormat(ssid)).Result()
	if err != nil {
		return false, err
	}
	return exists == 1, nil
}

func (Self *RedisJWTHandler) keyFormat(ssid string) string {
	return fmt.Sprintf("users:ssid:%s", ssid)
}
